tests/test_flow.py (1)

582-617: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Narrow the xfail scope so it doesn’t hide pw regressions.

xfail is currently applied to the whole test, so failures in the protein-water assertions are also treated as expected. That drops useful coverage beyond the known ww issue.

Suggested split to preserve `pw` coverage

-    `@pytest.mark.xfail`(
+    def test_batched_waters_have_protein_edges(self, batched_hetero_data):
+        """Ensure all waters in a batched graph have protein-water edges."""
+        updater = ProteinWaterUpdate(hidden_dims=(128, 16), layers=1)
+        edge_dict = updater.build_edges(batched_hetero_data, k_pw=4, k_ww=3)
+        pw_edges = edge_dict[("protein", "pw", "water")]
+        n_water = batched_hetero_data["water"].num_nodes
+        water_nodes_with_pw_edges = torch.unique(pw_edges[1])
+        assert len(water_nodes_with_pw_edges) == n_water, (
+            f"Only {len(water_nodes_with_pw_edges)}/{n_water} waters have protein edges in batched data"
+        )
+
+    `@pytest.mark.xfail`(
         reason=(
             "build_knn_edges' src/dst argument-order fix changes self-graph (ww) "
             "edge direction: row 0 now holds discovered neighbors rather than query "
             "points, so a point that is nobody's k-nearest neighbor can be dropped "
             "from coverage. The fixed-degree k_pw/k_ww KNN approach is replaced by "
             "radius-based edges + KNN-fallback-for-isolated-nodes in a future PR "
             "(edge type flags & dynamic edge construction), which removes the "
             "k_pw/k_ww params and fixes this guarantee structurally. will remove this "
             "marker when that PR is created."
         ),
         strict=True,
     )
-    def test_batched_waters_have_edges(self, batched_hetero_data):
-        """Ensure all waters in a batched graph have edges."""
+    def test_batched_waters_have_water_edges(self, batched_hetero_data):
+        """Ensure all waters in a batched graph have water-water edges."""
         updater = ProteinWaterUpdate(hidden_dims=(128, 16), layers=1)
-
         edge_dict = updater.build_edges(batched_hetero_data, k_pw=4, k_ww=3)
-        pw_edges = edge_dict[("protein", "pw", "water")]
         ww_edges = edge_dict[("water", "ww", "water")]
-
         n_water = batched_hetero_data["water"].num_nodes
-
-        # Check protein-water edges
-        water_nodes_with_pw_edges = torch.unique(pw_edges[1])
-        assert len(water_nodes_with_pw_edges) == n_water, (
-            f"Only {len(water_nodes_with_pw_edges)}/{n_water} waters have protein edges in batched data"
-        )
-
-        # Check water-water edges
         if n_water > 1:
             water_nodes_with_ww_edges = torch.unique(ww_edges[0])
             assert len(water_nodes_with_ww_edges) == n_water, (
                 f"Only {len(water_nodes_with_ww_edges)}/{n_water} waters have water-water edges in batched data"
             )

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/test_flow.py` around lines 582 - 617, The xfail marker is currently
applied to the entire test_batched_waters_have_edges function, which hides
failures in the protein-water edge assertions that should not be expected to
fail. Remove the xfail decorator from the function and instead apply it only to
the water-water edge checking section (the assertions checking
water_nodes_with_ww_edges). This can be done by either splitting the test into
two separate test functions with xfail only on the water-water test, or by
wrapping just the water-water edge assertion block with pytest.xfail() to
preserve protein-water edge coverage while still allowing the known water-water
edge failure.